import torch
import math


class AdamW(torch.optim.Optimizer):
    def __init__(
        self,
        params,
        lr: float = 1e-3,       # 学习率，默认为1e-3
        betas: tuple[float, float] = (0.9, 0.95),  # 用于计算梯度及其平方的移动平均值的系数
        eps: float = 1e-8,      # 添加到分母中以提高数值稳定性的小常数
        weight_decay: float = 0.1,  # 权重衰减系数，用于正则化
        **kwargs,
    ):
        # 设置默认参数字典
        defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay}
        # 调用父类构造函数初始化优化器
        super().__init__(params, defaults)

    @torch.no_grad()  # 禁用梯度计算，因为这是优化步骤，不需要梯度
    def step(self, closure=None):
        # 如果提供了闭包函数，则执行它并获取损失值
        loss = None if closure is None else closure()

        # 遍历所有参数组
        for group in self.param_groups:
            # 提取超参数并计算常数项，避免在循环中重复计算
            b1, b2 = group["betas"]  # 获取beta1和beta2
            lr = group["lr"]         # 获取学习率
            eps = group["eps"]       # 获取epsilon值
            wd = group["weight_decay"]  # 获取权重衰减系数
            # 预计算常数项，避免重复计算
            one_minus_b1 = 1 - b1    # 1 - beta1
            one_minus_b2 = 1 - b2    # 1 - beta2

            # 遍历当前参数组中的所有参数
            for p in group["params"]:
                # 如果参数没有梯度，则跳过
                if p.grad is None:
                    continue

                # 获取参数的状态字典（存储一阶矩、二阶矩和时间步）
                state = self.state[p]
                # 状态初始化，如果参数状态为空，则初始化
                if len(state) == 0:
                    # 初始化一阶矩估计为零张量
                    state["m"] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    # 初始化二阶矩估计为零张量
                    state["v"] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    # 初始化时间步为0
                    state["t"] = 0

                # 获取状态变量
                m, v, t = state["m"], state["v"], state["t"]
                t += 1  # 时间步加1
                grad = p.grad  # 获取当前参数的梯度

                # 使用in-place操作更新一阶矩和二阶矩估计
                # m = beta1 * m + (1 - beta1) * grad
                m.mul_(b1).add_(grad, alpha=one_minus_b1)
                # v = beta2 * v + (1 - beta2) * grad^2
                v.mul_(b2).addcmul_(grad, grad, value=one_minus_b2)

                # 计算偏差校正项，避免重复计算
                denom = 1 - b2 ** t      # 1 - beta2^t
                denom_sqrt = math.sqrt(denom)  # sqrt(1 - beta2^t)
                numer = 1 - b1 ** t      # 1 - beta1^t
                # 计算步长（包括偏差校正）
                step_size = lr * denom_sqrt / numer

                # 更新参数（in-place）
                # p = p - step_size * m / (sqrt(v) + eps)
                p.data.addcdiv_(m, v.sqrt().add_(eps), value=-step_size)

                # 应用权重衰减（AdamW特有的权重衰减方式）
                if wd != 0:
                    # p = p - lr * wd * p
                    p.data.add_(p.data, alpha=-lr * wd)

                # 保存更新后的时间步
                state["t"] = t

        return loss
